#ifndef __PA_VOCAB_H__
#define __PA_VOCAB_H__

#include <vector>
#include <unordered_map>
#include <string>
#include <map>

#include "common/common.h"
#include "model/model_base.h"

namespace pa {

class Vocab {
public:
    using Id = int32_t;
    using Token = std::string;

    // 两阶段分词
    // pre Tokenization
    // subword Tokenization
    enum VOCAB_PRE_TYPE {
        VOCAB_PRE_TYPE_QWEN2,
        VOCAB_PRE_TYPE_NONE
    }; // VOCAB_PRE_TYPE

    enum VOCAB_TYPE {
        VOCAB_TYPE_BPE,
        VOCAB_TYPE_UNK
    }; // VOCAB_TYPE

    enum VOCAB_POST_TYPE {

    }; // VOCAB_POST_TYPE

    Vocab(const ModelBase& model): m_model(model) {}
    ~Vocab() {}

    bool encoder_init();

    const Token& lookup(const Id& idx) const;
    Id lookup(const Token& token) const;

#ifndef PA_UNITTEST
private:
#endif // not define PA_UNITTEST

    bool engine_init_from_qwen_hub(const ModelBase& qwen2);

    std::unordered_map<Token, Id> m_encoder;
    std::vector<Token> m_decoder;

    std::unordered_map<Token, Id> m_added_token_encoder;
    std::vector<Token> m_added_token_decoder;

    const ModelBase& m_model;
}; // class Vocab
} // namespace pa

#endif // __PA_VOCAB_H__